#!/usr/bin/python
from pylab import *

# import numpy as np
# import matplotlib.pyplot as plt
# import matplotlib
from matplotlib.collections import PatchCollection
# import matplotlib.path as mpath
import matplotlib.patches as mpatches
import matplotlib.lines as mlines
import copy

# font = "sans-serif"
# fig = figure(figsize=(30,10))
# ax = fig.add_axes([0,0,1,1])  # the ticks not displayed
# ax = fig.add_axes([0.1,0.1,0.8,0.8])
# ax = fig.add_subplot(111)

fig, ax = plt.subplots()


# x = arange(0,1,0.01)
# plot(x,sin(2*pi*x))


# # create 3x3 grid to plot the artists
# # pos = np.mgrid[0.2:0.8:3j, 0.2:0.8:3j].reshape(2, -1)

# patches = []

# # add a circle
# # art = mpatches.Circle(pos[:,0], 0.1,ec="none")
# art = mpatches.Circle((0.5,0.5), 0.1,ec="none")
# patches.append(art)
# # ax.text(0.5, 0.5-0.15, "Circle", ha="center", family=font, size=14)

# art = mpatches.Rectangle((0.5,0.5), 0.5, 0.2, ec='none')
# # art = mpatches.Rectangle((0.5,0.5), 0.5, 0.2, fc='r', ec='none')
# patches.append(art)

# art = mpatches.Arrow(0,0,0.3,0.3,width=0.2)
# patches.append(art)

# # lines (not append to patches, added to ax)
# line = mlines.Line2D([0,0.2],[0.3,0.4], lw=4, color='r', alpha=0.4) # draw a line (0,0.3) - (0.2,0.4)
# lineX = linspace(0,2,5)
# lineY = [0,0,1,0,0]
# line = mlines.Line2D(lineX, lineY, lw=4, c='b', alpha=0.4)

# setp(line, color='g')           # set the properties of 'line'


# # colors = 100*np.random.rand(len(patches))
# collection = PatchCollection(patches, cmap=matplotlib.cm.jet, alpha=0.4)
# # collection.set_array(np.array(colors))
# # ax.add_collection(collection)


# ax.add_line(line)               # add lines
xlim(0,600)                       # must set the xlim/ylim manually
ylim(0,200)
# ax.set_xticks([])
# ax.set_yticks([])
ax.set_xticks([-1, -0.5, 0, 0.5, 1], ['$-\pi$', '$-\pi/2$', '$0$', '$+\pi/2$', '$+\pi$'])


# savefig("tmp.png", dpi=72)
# show()



#def plot_circle(center, radius, fill=True, color="black", **args):
# def plot_circle(center, radius, **args):
#     patch = mpatches.Circle(center, radius, kwargs=args)
#     ax.add_patch(patch)

 # def plot_circle(center, radius, **args):
 #    patch = mpatches.Circle(center, radius, kwargs=args)
 #    ax.add_patch(patch)
   
def plot_circle(center, radius, **args):
   patch = mpatches.Circle(center, radius, fill=args['fill'], color=args['color'])
   ax.add_patch(patch)

def plot_dot(pos):
    radius = 2
    plot_circle(pos, radius, fill=True, color="black")

    
def plot_rect(start_pos, end_pos, **args):
   # start_pos and end_pos any pair of the diagonal point
   if start_pos[0] < end_pos[0]:
      sx = start_pos[0]
      ex = end_pos[0]
   else:
      sx = end_pos[0]
      ex = start_pos[0]

   if start_pos[1] < end_pos[1]:
      sy = start_pos[1]
      ey = end_pos[1]
   else:
      sy = end_pos[1]
      ey = start_pos[1]

   width = ex - sx + 1
   height = ey - sy + 1
   
   # sx = start_pos[0] if start_pos[0]<end_pos[0] else end_pos[0]
   # sy = start_pos[1] if start_pos[1]<end_pos[1] else end_pos[1]
   
   patch = mpatches.Rectangle((sx, sy), width, height, fill=args['fill'], color=args['color'])
   ax.add_patch(patch)


# ==================== parameters and global vars ====================
width = 8
# width = 16                      # 16x16 multiply
# width = 32

col_middle = width
pps_num = width/2+1                      # num of pps
pps_width = width + 1                    # the width of last PP is width
cols_num = width*2
rows_num = width/2+1

# dots[rows_num][cols_num]
dots = [[None]*cols_num for i in range(rows_num)]


step_x = 10                      # the diameter of a node in the grid
step_y = 10
# pos_base = (20,20)

# def display_dots(pos_base=(0,0), step_x=10, step_y=10):
def display_dots(dots, pos_base=(0,0)):
   point_base = (pos_base[0]+step_x/2, pos_base[1]+step_y/2)
   # for y in range(rows_num):
   #     for x in range(cols_num):
   #         pos = (pos_base[0]+step_x*(cols_num-x), pos_base[1]+step_y*(rows_num-y))
   #         if dots[y][x] == None:
   #             continue
   #         elif dots[y][x] == 1:
   #             plot_circle(pos, radius)
   #         else:
   #             text(pos[0], pos[1], dots[y][x], ha="center", va="center", family="sans-serif", size=10)

   # add one row for display ticks
   disp_rows_num = rows_num + 1
   ticks = [str(i) for i in range(cols_num)]
   disp_dots = [ticks] + dots
   for y in range(disp_rows_num):
      for x in range(cols_num):
         pos = (point_base[0]+step_x*(cols_num-1-x), point_base[1]+step_y*(disp_rows_num-1-y))
         if disp_dots[y][x] == None:
            continue
         elif disp_dots[y][x] == 1:
            plot_dot(pos)
         else:
            text(pos[0], pos[1], disp_dots[y][x], ha="center", va="center", family="sans-serif", size=10)

   # draw the grid
   pos_left_buttom = (pos_base[0], pos_base[1])
   pos_right_top = (pos_base[0]+step_x*cols_num, pos_base[1]+step_y*disp_rows_num)
   # the boundaries
   plot_rect(pos_left_buttom, pos_right_top, fill=False, color='black')
   # the table header
   tab_header_start_pos = (pos_left_buttom[0], pos_right_top[1]-step_y)
   tab_header_end_pos = pos_right_top
   plot_rect(tab_header_start_pos, tab_header_end_pos, fill=True, color='red', alpha=0.8)

   # the columns
   # lineX = [pos_base[0]+step_x*i for i in range(1,cols_num)]
   for i in range(1,cols_num):
      x = pos_base[0]+step_x*i
      y0 = pos_left_buttom[1]
      y1 = pos_right_top[1]
      line = mlines.Line2D([x,x], [y0, y1], lw=1, c='black')
      ax.add_line(line)
   # the rows
   for i in range(1, disp_rows_num):
      y = pos_base[1]+step_y*i
      x0 = pos_left_buttom[0]
      x1 = pos_right_top[0]
      line = mlines.Line2D([x0,x1], [y,y], lw=1, c='black')
      ax.add_line(line)

      
# the dots (PPs)
for row in range(rows_num):
    start_x = row*2
    if row == rows_num-1:       # the last PP has one column less than others
        end_x = start_x + pps_width-1
    else:
        end_x = start_x + pps_width
    for col in range(start_x, end_x):
        dots[row][col] = 1   # display the dot
# the S[]
for col in range(0, col_middle, 2):
    dots[1+col/2][col] = "S[%d]" % (col/2)
# the E[] & 1
for col in range(col_middle+3, cols_num-1+2, 2):
    row = (col-col_middle-3)/2+1
    dots[row][col] = "E[%d]" % row
    if col < cols_num-1:
        dots[row][col+1] = "1"
dots[0][col_middle+1] = "{~E[0]}"
dots[0][col_middle+2] = "{~E[0]}"
dots[0][col_middle+3] = "E[0]"


# calc the items of each column
cols = []
for col in range(cols_num):
   items = []
   for row in range(rows_num):
      if dots[row][col] == None:
         continue
      elif dots[row][col] == 1:
         items.append("pp[%d][%d]" % (row, col-2*row))
      else:
         items.append(dots[row][col])
   cols.append(items)
#   print "cols[%d] = %s;" % (col, ' + '.join(items))

display_dots(dots, pos_base=(0,100))
# print "cols = ", cols


# pull up the dots
dots_up = [[None]*cols_num for i in range(rows_num)]
for col in range(cols_num):
   items = []
   for row in range(rows_num):
      if dots[row][col]!= None:
         items.append(dots[row][col])
   items += [None]*(rows_num - len(items))
   for row in range(rows_num):
      dots_up[row][col] = items[row]

display_dots(dots_up, pos_base=(0,0))

# just treat all as dots, no text values
cols_len = []
for col in range(cols_num):
   cols_len.append(len(cols[col]))
# cols_carry = [0] * cols_num

# ==================== wallace tree ====================
# def wallace_compress():
#    global cols_len
#    compress_fields = []
#    cols_carry[-1] = 0              # for column 0
#    for col in range(cols_num):
#       n_compress_3_2 = cols_len[col] / 3
#       n_reminder = cols_len[col] % 3
# #      if n_compress_3_2 == 0 and n_reminder+
#       # calc the length after compress (next stage)
#       cols_len[col] = n_compress_3_2 + (n_reminder>0) + cols_carry[col-1]
#       cols_carry[col] = n_compress_3_2 + (n_reminder==2)
#       # calc the compress fields for current stage
#       fields = [3]*n_compress_3_2
#       if n_reminder == 2:
#          fields.append(2)
#       elif n_reminder == 1:
#          fields.append(1)
#       compress_fields.append(fields)
#    return compress_fields
         
         
##################################################
compress_array = []
compress_stage = len(compress_array) - 1

def wallace_compress():
   global cols, compress_stage, compress_array
   
   compress_array.append([[] for i in range(cols_num)])
   compress_stage = len(compress_array) - 1
   
   compress_fields = []
   cols_carry = [[] for i in range(cols_num)]

   for col in range(cols_num):
      n_compress_3_2 = len(cols[col]) / 3
      n_reminder = len(cols[col]) % 3
      
      curr_col_result = []
      fields = []
      # clac current col's sum
      i = -1
      for i in range(n_compress_3_2):
         comp_in = cols[col][i*3:(i+1)*3]       # from i*3 to i*3+2
         comp_sum = "S%d_%d_%d" % (compress_stage+1, col, i) # stage_col_row
         comp_carry = "S%d_%d_Carry_%d" % (compress_stage+1, col+1, i)  # this is not the real name which is determined by next col
         compress_io = [comp_carry, comp_sum] + comp_in # (3,2) counter / full adder
         
         curr_col_result.append(comp_sum)
         cols_carry[col].append(comp_carry)
         compress_array[compress_stage][col].append(compress_io)

         fields.append(3)
      
      i += 1
      if n_reminder == 2:
         comp_in = cols[col][i*3:]
         assert(len(comp_in) == 2)
         comp_sum = "S%d_%d_%d" % (compress_stage+1, col, i)
         comp_carry = "S%d_%d_Carry_%d" % (compress_stage+1, col+1, i)
         compress_io = [comp_carry, comp_sum] + comp_in # (2,2) counter/half adder

         curr_col_result.append(comp_sum)
         cols_carry[col].append(comp_carry)
         compress_array[compress_stage][col].append(compress_io)

         fields.append(2)

      elif n_reminder == 1:
         comp_in = cols[col][i*3]
         comp_sum = "S%d_%d_%d" % (compress_stage+1, col, i) 
         compress_io = [comp_sum, comp_in]              # just a wire
         
         curr_col_result.append(comp_sum)
         compress_array[compress_stage][col].append(compress_io)
      
         fields.append(1)

      compress_fields.append(fields)

      # calc the previous col's carry out (get its real name)
      for j in range(len(cols_carry[col-1])):
         prev_carry = "S%d_%d_%d" % (compress_stage+1, col, i+j+1)
         curr_col_result.append(prev_carry)
         cols_carry[col-1][j] = prev_carry # rename to the real name
         compress_array[compress_stage][col-1][j][0] = prev_carry # rename the previous col's carry out 

      cols[col] = curr_col_result
   return compress_fields
##################################################


# draw oval vertical
def plot_oval(start_pos, end_pos, **args):
   # patch = mpatches.FancyBboxPatch((pos[0]-step_x/2,pos[1]-step_y/2), step_x, 30, fill=False, color="blue", linewidth=3, boxstyle=mpatches.BoxStyle("Round", pad=0))
   if start_pos[0] < end_pos[0]:
      sx = start_pos[0]
      ex = end_pos[0]
   else:
      sx = end_pos[0]
      ex = start_pos[0]

   if start_pos[1] < end_pos[1]:
      sy = start_pos[1]
      ey = end_pos[1]
   else:
      sy = end_pos[1]
      ey = start_pos[1]

   width = ex - sx + 1
   height = ey - sy + 1

   patch = mpatches.FancyBboxPatch((sx, sy), width, height, fill=args['fill'], color=args['color'], linewidth=2, boxstyle=mpatches.BoxStyle("Round", pad=3))
   ax.add_patch(patch)

   


def dot2pos(row, col, pos_base):
   disp_rows_num = rows_num + 1
   point_base = (pos_base[0]+step_x/2, pos_base[1]+step_y/2)
   pos = (point_base[0]+step_x*(cols_num-1-col), point_base[1]+step_y*(disp_rows_num-1-(row+1)))
   return pos
      
      


# def display_compress(cols_length, compress_fields, pos_base=(0,0)):
#    dots_compress = [[None]*cols_num for i in range(rows_num)]
#    for col in range(cols_num):
#       # items = [1] * cols_length[col] + [None] * (rows_num - cols_length[col])
#       for row in range(cols_length[col]):
#          dots_compress[row][col] = 1                    # valid dot
#    display_dots(dots_compress, pos_base)
#    # draw the select oval
#    for col in range(cols_num):
#       cnt = 0
#       for n in compress_fields[col]:
#          if n==2 or n==3:
#             spos, epos = dot2pos(cnt, col, pos_base), dot2pos(cnt+n-1, col, pos_base)
#             plot_oval(spos, epos, color="red", fill=False)
#          cnt += n

def display_compress(cols_val, compress_fields, pos_base=(0,0)):
   dots_compress = [[None]*cols_num for i in range(rows_num)]
   for col in range(cols_num):
      for row in range(len(cols_val[col])):
         dots_compress[row][col] = 1                    # valid dot
   display_dots(dots_compress, pos_base)
   
   if compress_fields == None:
      return
   # draw the select oval
   for col in range(cols_num):
      cnt = 0
      for n in compress_fields[col]:
         if n==2 or n==3:
            spos, epos = dot2pos(cnt, col, pos_base), dot2pos(cnt+n-1, col, pos_base)
            plot_oval(spos, epos, color="red", fill=False)
         cnt += n

# print cols_len

# cols_len_s0 = copy.deepcopy(cols_len)
# compress_fields_s0 = wallace_compress()
# display_compress(cols_len_s0, compress_fields_s0, pos_base=(200,0))

# cols_len_s1 = copy.deepcopy(cols_len)
# compress_fields_s1 = wallace_compress()
# display_compress(cols_len_s1, compress_fields_s1, pos_base=(200,100))

# cols_len_s2 = copy.deepcopy(cols_len)
# compress_fields_s2 = wallace_compress()
# display_compress(cols_len_s2, compress_fields_s2, pos_base=(400,0))

# cols_len_s3 = copy.deepcopy(cols_len)
# compress_fields_s3 = wallace_compress()
# display_compress(cols_len_s3, compress_fields_s3, pos_base=(400,100))

# print cols_len

cols_orig = copy.deepcopy(cols)

cnt = 0
while True:
   cols_sx = copy.deepcopy(cols)
   print "stage%d : " % cnt, cols_sx
   compress_fields_sx = wallace_compress()
   display_compress(cols_sx, compress_fields_sx, pos_base=(200+(cnt/2)*200,100-(cnt%2)*100 ))
   # display_compress(cols_len_sx, compress_fields_sx, pos_base=(200+(cnt/2)*400,100-(cnt%2)*200 ))
   cnt += 1
   if max([len(i) for i in cols]) == 2:
      display_compress(cols, None, pos_base=(200+(cnt/2)*200,100-(cnt%2)*100 ))
      break
   

print "need %d stages to compress" % cnt

# the final 2 PPs
final_pp0 = ""
final_pp1 = ""
for col in range(cols_num-1,-1,-1):
   final_pp0 += cols[col][0] + ","
   if len(cols[col]) == 2:
      final_pp1 += cols[col][1] + ","
   elif len(cols[col]) == 1:
      final_pp1 += "1'b0,"
   else:
      assert(False)
final_pp0 = "{" + final_pp0[:-1] + "}"
final_pp1 = "{" + final_pp1[:-1] + "}"

print "==================== the final 2 PPs ===================="
print final_pp0
print final_pp1

# gen code for compress
compress_operations = [[] for i in range(cnt)]
create_wires = [[] for i in range(cnt)]
FA_cnt = 1
HA_cnt = 1
# for stage in range(cnt):
#    # compress_array[stage][col]
#    for col in range(cols_num):
#       for item in compress_array[stage][col]:
#          if len(item) == 5:     # full adder
#             compress_operations[stage].append("FA fa%d (%s);" % (FA_cnt, ", ".join(item)))
#             create_wires[stage] += item[0:1]
#             FA_cnt += 1
#          elif len(item) == 4:   # half adder
#             compress_operations[stage].append("HA ha%d (%s);" % (HA_cnt, ", ".join(item)))
#             create_wires[stage] += item[0:1]
#             HA_cnt += 1
#          elif len(item) == 2:   # just wire
#             compress_operations[stage].append("assign %s = %s;" % (item[0], item[1]))
#             create_wires[stage] += [item[0]]
#          else:
#             assert(False)

for stage in range(cnt):
   for col in range(cols_num):
      for item in compress_array[stage][col]:
         tmp_str = ""
         for i in item:
            if i == '1':        # replace "1" width "1'b1"
               tmp_str += "1'b1,"
            else:
               tmp_str += i+","
         tmp_str = tmp_str[:-1]

         if len(item) == 5:     # full adder
            compress_operations[stage].append("FA fa%d (%s);" % (FA_cnt, tmp_str))
            create_wires[stage] += item[0:1]
            FA_cnt += 1
         elif len(item) == 4:   # half adder
            compress_operations[stage].append("HA ha%d (%s);" % (HA_cnt, tmp_str))
            create_wires[stage] += item[0:1]
            HA_cnt += 1
         elif len(item) == 2:   # just wire
            # compress_operations[stage].append("assign %s = %s;" % (item[0], tmp_str))
            compress_operations[stage].append("assign %s;" % (tmp_str.replace(",", " = ")))
            create_wires[stage] += [item[0]]
         else:
            assert(False)




wires_str = ""
print "// ==================== wires need create ===================="
for stage in range(cnt):
   wires_str += "logic " + ','.join(create_wires[stage]) + ";\n"
print wires_str

# compress_str = ""
compress_str = wires_str
# compress_str += "logic [%d:0] final_pp0, final_pp1;" % (cols_num-1)
for stage in range(cnt):
   compress_str += "// ==================== Stage%d ====================" % (stage+1)
   compress_str += '\n'
   for op in compress_operations[stage]:
      compress_str += op + '\n'

compress_str += "// ==================== Final 2 PPs ====================\n"
compress_str += "assign final_pp0 = %s;\n" % final_pp0
compress_str += "assign final_pp1 = %s;\n" % final_pp1

print compress_str   
print "-"*50
print "total %d FAs and %d HAs" % (FA_cnt-1, HA_cnt-1)




# replace the 
f = open("booth2_multiply_template.sv", "r")
contents = f.read()
f.close()

contents = contents.replace("$width", str(width))
contents = contents.replace("$compress_code;", compress_str)

f = open("mybooth2.sv", "w")
f.write(contents)
f.close()



for i,v  in enumerate(cols_orig):
   print "cols[%d] = %s;" % (i, ' + '.join(v))




pos = dot2pos(1, 2, pos_base=(0,0))
plot_circle(pos, 8, fill=False, color="black")
   

# for i in range(16):
#     plot_circle((i*radius*3, 2), radius)

pos1 = dot2pos(1, 3, pos_base=(0,0))
pos2 = dot2pos(1+1, 3, pos_base=(0,0))
plot_oval(pos1, pos2, color="red", fill=False)


# patch1 = mpatches.FancyBboxPatch((pos[0]-step_x/2,pos[1]-step_y/2), step_x, 30, fill=False, color="blue", linewidth=2, boxstyle=mpatches.BoxStyle("Round", pad=0))
# patch2 = mpatches.Rectangle((pos[0]-step_x/2, pos[1]-step_y/2), step_x, 50, fill=False, color='red')
# ax.add_patch(patch1)
# ax.add_patch(patch2)


savefig("tmp.png", dpi=72)
show()
